{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c149523",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "from typing import *\n",
    "import time\n",
    "import copy\n",
    "import warnings\n",
    "\n",
    "from data.openai import *\n",
    "from data.generation import *\n",
    "from data.finetune import *\n",
    "from data.inference import *\n",
    "from data.io import *\n",
    "from data.evaluation import *\n",
    "from data.split import *\n",
    "\n",
    "from utils.paths import *\n",
    "from utils.metadata import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7ad01d7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import openai\n",
    "openai.api_key = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6264f98",
   "metadata": {},
   "outputs": [],
   "source": [
    "ALL_DATASETS = [\n",
    "    \"single_eq\",\n",
    "    \"addsub\",\n",
    "    \"multiarith\",\n",
    "    \"gsm8k\",\n",
    "    \"aqua\",\n",
    "    \"svamp\",\n",
    "    \n",
    "    \"date_understanding\",\n",
    "    \"coin_flip\",\n",
    "    \n",
    "    \"tracking_shuffled_objects\",\n",
    "    \"last_letter_concatenation\",\n",
    "    \n",
    "    \"commonsense_qa\",\n",
    "    \"strategy_qa\",\n",
    "]\n",
    "datasets = ALL_DATASETS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "477cf799",
   "metadata": {},
   "source": [
    "# Teacher Zero-shot-CoT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e52f09",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for dataset_key in datasets:\n",
    "    print(\" {} \".format(dataset_key).center(80, \"#\"))\n",
    "    completion_key = \"zs_cot\"\n",
    "    model_key = \"text-davinci-002\"\n",
    "    train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "    all_indices = train_indices + test_indices  # subset of aqua, gsm8k, for which we subsampled 10000 for train\n",
    "    generate_cot_completions(completion_key, dataset_key, model_key, indices=all_indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "885aa795",
   "metadata": {},
   "source": [
    "# Student Zero-Shot-CoT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f7e5450",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for dataset_key in datasets:\n",
    "        print(\"#\" * 80)\n",
    "        print(\"Inferring test zs_cot for {}-{}\".format(model_key, dataset_key))\n",
    "        print(\"#\" * 80)\n",
    "        completion_key = \"zs_cot\"\n",
    "        dataset = load_dataset(dataset_key)\n",
    "        train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "        generate_cot_completions(completion_key, dataset_key, model_key, indices=test_indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d172c5c",
   "metadata": {},
   "source": [
    "# Student Zero-Shot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbcd3ba2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "datasets = ALL_DATASETS\n",
    "for model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for dataset_key in datasets:\n",
    "        print(\"#\" * 80)\n",
    "        print(\"Inferring test zs for {}-{}\".format(model_key, dataset_key))\n",
    "        print(\"#\" * 80)\n",
    "        completion_key = \"zs\"\n",
    "        dataset = load_dataset(dataset_key)\n",
    "        train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "        infer_cot_completions(completion_key, dataset_key, model_key, template=None, split=\"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0fc2a35",
   "metadata": {},
   "source": [
    "# Student Few-Shot-CoT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c85db7ba",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "datasets = ALL_DATASETS\n",
    "for model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for dataset_key in datasets:\n",
    "        if dataset_key == \"tracking_shuffled_objects\":\n",
    "            continue\n",
    "        print(\"#\" * 80)\n",
    "        print(\"Inferring test fs-cot for {}-{}\".format(model_key, dataset_key))\n",
    "        print(\"#\" * 80)\n",
    "        completion_key = \"fs_cot_long\"\n",
    "        dataset = load_dataset(dataset_key)\n",
    "        infer_cot_completions(completion_key, dataset_key, model_key, template=\"few_shot_cot\", split=\"test\", max_tokens=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5e23b12",
   "metadata": {},
   "source": [
    "# Curate and Fine-Tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25bde0df",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for dataset_key in datasets:\n",
    "    train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "\n",
    "    # Get train completions\n",
    "    completion_data = load_completion_data(\"zs_cot\", dataset_key, \"text-davinci-002\")\n",
    "\n",
    "    # Generate FT data\n",
    "    template = \"special\"\n",
    "    file_key = \"zs_cot_{}_{}_train\".format(template, dataset_key)\n",
    "    generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                          indices=train_indices)\n",
    "\n",
    "    # Create file on OpenAI\n",
    "    create_finetune_file(file_key)  # openai\n",
    "    \n",
    "    # Create finetunes\n",
    "    for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        model_key = \"{}_{}\".format(base_model, file_key)\n",
    "        create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5992d13c",
   "metadata": {},
   "source": [
    "# Fetch Fine-tune-CoT Student Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44fd565f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for dataset_key in datasets:\n",
    "        total += 1\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_train\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        if get_model_id(model_key, strict=False):\n",
    "            print(\"{:60s} exists\".format(model_key))\n",
    "            complete += 1\n",
    "            continue\n",
    "\n",
    "        finetune_id = get_finetune_id(model_key)\n",
    "        if finetune_id is None:\n",
    "            print(\"{:60s} no finetune found\".format(model_key))\n",
    "            continue\n",
    "\n",
    "        response = openai.FineTune.retrieve(finetune_id)\n",
    "        model_id = response[\"fine_tuned_model\"]\n",
    "        if model_id is not None:\n",
    "            set_model_id(model_key, model_id)\n",
    "            print(\"{:60s} fetched\".format(model_key))\n",
    "            print(\"    {:40s}\".format(model_id))\n",
    "            complete += 1\n",
    "        else:\n",
    "            print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "\n",
    "\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "288ce254",
   "metadata": {},
   "source": [
    "# Fine-Tune-CoT Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d121231",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot\"\n",
    "for dataset_key in datasets:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        print(\"#\" * 80)\n",
    "        print(dataset_key, base_model_key)\n",
    "        print(\"#\" * 80)\n",
    "        file_key = \"zs_cot_{}_{}_train\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        _ = infer_cot_completions(completion_key, dataset_key, model_key, template=template)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e419f2f8",
   "metadata": {},
   "source": [
    "### Long Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d662df8d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot_long\"\n",
    "for dataset_key in datasets:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        print(\"#\" * 80)\n",
    "        print(dataset_key, base_model_key)\n",
    "        print(\"#\" * 80)\n",
    "        file_key = \"zs_cot_{}_{}_train\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        _ = infer_cot_completions(completion_key, dataset_key, model_key, template=template, max_tokens=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f05d44a",
   "metadata": {},
   "source": [
    "# Fine-tune-CoT (Other Teachers)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fad80f1c",
   "metadata": {},
   "source": [
    "### Teacher Zero-shot-CoT "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba69e2ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ablation_datasets = [\"multiarith\", \"svamp\", \"date_understanding\", \"last_letter_concatenation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24ad7b12",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for dataset_key in model_ablation_datasets:\n",
    "    print(\" {} \".format(dataset_key).center(80, \"#\"))\n",
    "    completion_key = \"zs_cot\"\n",
    "    for model_key in [\"davinci\", \"text-davinci-001\", \"text-davinci-003\"]:\n",
    "        train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "        all_indices = train_indices + test_indices  # subset of aqua, gsm8k, for which we subsampled 10000 for train\n",
    "        generate_cot_completions(completion_key, dataset_key, model_key, indices=all_indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a4ebaf6",
   "metadata": {},
   "source": [
    "### Fine-tune Students"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e24876fc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for dataset_key in model_ablation_datasets:\n",
    "    for teacher_model in [\"davinci\", \"text-davinci-001\", \"text-davinci-003\"]:\n",
    "        train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "\n",
    "        # Get train completions\n",
    "        completion_data = load_completion_data(\"zs_cot\", dataset_key, teacher_model)\n",
    "        \n",
    "        evaluation = evaluate_completions(completion_data, dataset_key, template=None, indices=train_indices)\n",
    "        if evaluation.correct.sum() == 0:\n",
    "            print(\"No correct samples for {:20s} {:20s}\".format(dataset_key, teacher_model))\n",
    "            continue\n",
    "\n",
    "        # Generate FT data\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_{}_train\".format(template, dataset_key, teacher_model)  # third key added\n",
    "        generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                              indices=train_indices)\n",
    "\n",
    "        # Create file on OpenAI\n",
    "        create_finetune_file(file_key)  # openai\n",
    "\n",
    "        # Create finetunes\n",
    "        for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            model_key = \"{}_{}\".format(base_model, file_key)\n",
    "            create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89c903bf",
   "metadata": {},
   "source": [
    "### Fetch Students"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71d0dd8a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "for teacher_model in [\"davinci\", \"text-davinci-001\", \"text-davinci-003\"]:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        for dataset_key in model_ablation_datasets:\n",
    "            total += 1\n",
    "            template = \"special\"\n",
    "            file_key = \"zs_cot_{}_{}_{}_train\".format(template, dataset_key, teacher_model)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False):\n",
    "                print(\"{:80s} exists\".format(model_key))\n",
    "                complete += 1\n",
    "                continue\n",
    "\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id is None:\n",
    "                print(\"{:80s} no finetune found\".format(model_key))\n",
    "                continue\n",
    "\n",
    "            response = openai.FineTune.retrieve(finetune_id)\n",
    "            model_id = response[\"fine_tuned_model\"]\n",
    "            if model_id is not None:\n",
    "                set_model_id(model_key, model_id)\n",
    "                print(\"{:80s} fetched\".format(model_key))\n",
    "                print(\"    {:40s}\".format(model_id))\n",
    "                complete += 1\n",
    "            else:\n",
    "                print(\"{:80s} {}\".format(model_key, response[\"status\"]))\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "430681d0",
   "metadata": {},
   "source": [
    "### Run Students"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d62d42",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot\"\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for teacher_model in [\"davinci\", \"text-davinci-001\", \"text-davinci-003\"]:\n",
    "        for dataset_key in model_ablation_datasets:\n",
    "            print(\"#\" * 80)\n",
    "            print(\"{}-{}-{}\".format(base_model_key, teacher_model, dataset_key))\n",
    "            print(\"#\" * 80)\n",
    "            file_key = \"zs_cot_{}_{}_{}_train\".format(template, dataset_key, teacher_model)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False):\n",
    "                _ = infer_cot_completions(completion_key, dataset_key, model_key, template=template)\n",
    "            else:\n",
    "                print(\"Skipping unavailable model {}\".format(model_key))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5c4b9fb",
   "metadata": {},
   "source": [
    "### Run Students (Long Inference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c74be93",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot_long\"\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for teacher_model in [\"davinci\", \"text-davinci-001\", \"text-davinci-003\"]:\n",
    "        for dataset_key in model_ablation_datasets:\n",
    "            print(\"#\" * 80)\n",
    "            print(\"{}-{}-{}\".format(base_model_key, teacher_model, dataset_key))\n",
    "            print(\"#\" * 80)\n",
    "            file_key = \"zs_cot_{}_{}_{}_train\".format(template, dataset_key, teacher_model)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False):\n",
    "                _ = infer_cot_completions(completion_key, dataset_key, model_key, template=template, max_tokens=1024)\n",
    "            else:\n",
    "                print(\"Skipping unavailable model {}\".format(model_key))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32f491fd",
   "metadata": {},
   "source": [
    "# ☆ Fine-Tune-CoT w/ Template Splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9bb85cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_datasets = [\"multiarith\", \"date_understanding\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5d2d56b",
   "metadata": {},
   "source": [
    "### Teacher inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ccde47",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for dataset_key in ts_datasets:\n",
    "    dataset = load_dataset(dataset_key)\n",
    "    train_indices, test_indices = get_train_test_indices(dataset_key, split_key=\"template_split\")\n",
    "\n",
    "    # Get train completions\n",
    "    completion_data = load_completion_data(\"zs_cot\", dataset_key, \"text-davinci-002\")\n",
    "    train_completions = {i: completion_data[i] for i in train_indices}\n",
    "    test_completions = {i: completion_data[i] for i in test_indices}\n",
    "\n",
    "    # Generate FT data\n",
    "    template = \"special\"\n",
    "    file_key = \"zs_cot_{}_{}_template_train\".format(template, dataset_key)\n",
    "    generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                          indices=train_indices)\n",
    "\n",
    "    # Create file on OpenAI\n",
    "    create_finetune_file(file_key)  # openai\n",
    "    \n",
    "    # Create finetunes\n",
    "    for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        model_key = \"{}_{}\".format(base_model, file_key)\n",
    "        create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac60d106",
   "metadata": {},
   "source": [
    "### Fetch Students"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d930644",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for dataset_key in ts_datasets:\n",
    "        total += 1\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_template_train\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        if get_model_id(model_key, strict=False):\n",
    "            print(\"{:60s} exists\".format(model_key))\n",
    "            complete += 1\n",
    "            continue\n",
    "\n",
    "        finetune_id = get_finetune_id(model_key)\n",
    "        if finetune_id is None:\n",
    "            print(\"{:60s} no finetune found\".format(model_key))\n",
    "            continue\n",
    "\n",
    "        response = openai.FineTune.retrieve(finetune_id)\n",
    "        model_id = response[\"fine_tuned_model\"]\n",
    "        if model_id is not None:\n",
    "            set_model_id(model_key, model_id)\n",
    "            print(\"{:60s} fetched\".format(model_key))\n",
    "            print(\"    {:40s}\".format(model_id))\n",
    "            complete += 1\n",
    "        else:\n",
    "            print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "\n",
    "\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f3da67a",
   "metadata": {},
   "source": [
    "###  Student Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b0aa64e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot_long\"\n",
    "for dataset_key in ts_datasets:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        print(\"#\" * 80)\n",
    "        print(dataset_key, base_model_key)\n",
    "        print(\"#\" * 80)\n",
    "        file_key = \"zs_cot_{}_{}_template_train\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        train, test = get_train_test_indices(dataset_key, split_key=\"template_split\")\n",
    "        _ = infer_cot_completions(completion_key, dataset_key, model_key, indices=test, template=template, max_tokens=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1573ab9b",
   "metadata": {},
   "source": [
    "# ☆ Fine-Tune-CoT w/ Golden Filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "377f6465",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_key = \"date_understanding\"\n",
    "completions = load_completion_data(\"zs_cot\", dataset_key, \"text-davinci-002\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93fc00f4",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "train, test = get_train_test_indices(dataset_key)\n",
    "for i in train:\n",
    "    print(\"Sample #{:03d}\".format(i))\n",
    "    q = completions[i][0][\"question\"]\n",
    "    a = completions[i][0][\"answer\"]\n",
    "    r = completions[i][0][\"reasoning_completion\"]\n",
    "    c = completions[i][0][\"completion\"]\n",
    "    \n",
    "    clean_a = cleanse_answer(a, dataset_key)\n",
    "    clean_prediction, prediction_candidates = cleanse_prediction(c, dataset_key, answer_prefix=None)\n",
    "    \n",
    "    if clean_a == clean_prediction:\n",
    "        print(\"-\" * 80)\n",
    "        print(q.strip())\n",
    "        print(\"-\" * 80)\n",
    "        print(r.strip())\n",
    "        print(c.strip())\n",
    "        print(\"-\" * 80)\n",
    "    else:\n",
    "        print(\"-\" * 80)\n",
    "        print(\"WRONG\")\n",
    "        print(\"-\" * 80)\n",
    "    \n",
    "    break  # for human-filtering..........."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fb232a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "train, test = get_train_test_indices(dataset_key)\n",
    "completion_data = load_completion_data(\"zs_cot\", dataset_key, \"text-davinci-002\")\n",
    "\n",
    "correct = []\n",
    "for i in train:\n",
    "    sample = completion_data[i][0]\n",
    "    answer_prefix = \"-->\"\n",
    "    prediction = cleanse_prediction(sample[\"completion\"], dataset_key, answer_prefix=answer_prefix,\n",
    "                                    return_all=False)\n",
    "    answer = cleanse_answer(sample[\"answer\"], dataset_key)\n",
    "    is_correct_completion = compare_prediction_and_answer(prediction, answer, dataset_key)\n",
    "    if is_correct_completion:\n",
    "        correct.append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71447a3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(correct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ee3aba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "good = [6,8,10,16,18,19,27,29,30,35,36,44,45,46,59,61,63,67,68,71,75,76,78,79,80,90,97,112,114,122,133,134,136,138,139,140,141,142,144,145,146,149,152,154,162,164,166,168,170,179,182,184,186,188,190,191,194,196,199,204,208,210,213,215,216,223,226,228,229,231,232,234,236,238,239,241,246,252,253,254,255,258,259,260,261,263,266,268,269,270,271,272,275,276,278,284,287,295,296,297,298,299,300,301,302,304,305,306,307,308,309,311,316,319,320,324,330,334,354,356,357,360,368]\n",
    "bad = [5,7,11,20,24,33,34,37,52,54,64,65,86,89,92,101,124,126,135,155,157,158,159,160,171,198,200,205,206,212,218,230,233,245,248,249,281,282,283,286,293,313,329,344,350,352,363]\n",
    "\n",
    "assert set(good).intersection(set(bad)) == set()\n",
    "assert set(good).issubset(set(correct))\n",
    "assert set(bad).issubset(set(correct))\n",
    "assert set(good + bad) == set(correct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1810defe",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(train), len(correct), len(good), len(bad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f984652",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = np.random.RandomState(0)\n",
    "good_baseline = state.permutation(correct)[:len(good)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9f5c75e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for dataset_key in [\"date_understanding\"]:\n",
    "    dataset = load_dataset(dataset_key)\n",
    "\n",
    "    # Get train completions\n",
    "    completion_data = load_completion_data(\"zs_cot\", dataset_key, \"text-davinci-002\")\n",
    "    \n",
    "    # Generate FT data\n",
    "    template = \"special\"\n",
    "    file_key = \"zs_cot_{}_{}_good_baseline\".format(template, dataset_key)\n",
    "    generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                          indices=good_baseline)\n",
    "\n",
    "    # Create file on OpenAI\n",
    "    create_finetune_file(file_key)  # openai\n",
    "    \n",
    "    # Create finetunes\n",
    "    for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        model_key = \"{}_{}\".format(base_model, file_key)\n",
    "        create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbbefd55",
   "metadata": {},
   "source": [
    "### Fetch Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca67453f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for dataset_key in [\"date_understanding\"]:\n",
    "        total += 1\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_good_baseline\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        if get_model_id(model_key, strict=False):\n",
    "            print(\"{:60s} exists\".format(model_key))\n",
    "            complete += 1\n",
    "            continue\n",
    "\n",
    "        finetune_id = get_finetune_id(model_key)\n",
    "        if finetune_id is None:\n",
    "            print(\"{:60s} no finetune found\".format(model_key))\n",
    "            continue\n",
    "\n",
    "        response = openai.FineTune.retrieve(finetune_id)\n",
    "        model_id = response[\"fine_tuned_model\"]\n",
    "        if model_id is not None:\n",
    "            set_model_id(model_key, model_id)\n",
    "            print(\"{:60s} fetched\".format(model_key))\n",
    "            print(\"    {:40s}\".format(model_id))\n",
    "            complete += 1\n",
    "        else:\n",
    "            print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "\n",
    "\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad3e256b",
   "metadata": {},
   "source": [
    "###  Student Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5c99a42",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot_long\"\n",
    "for dataset_key in [\"date_understanding\"]:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        print(\"#\" * 80)\n",
    "        print(dataset_key, base_model_key)\n",
    "        print(\"#\" * 80)\n",
    "        file_key = \"zs_cot_{}_{}_good\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        train, test = get_train_test_indices(dataset_key)\n",
    "        _ = infer_cot_completions(completion_key, dataset_key, model_key, indices=test, template=template, max_tokens=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34da9dac",
   "metadata": {},
   "source": [
    "# ☆ Fine-Tune-CoT w/ All Samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "413269de",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for dataset_key in [\"date_understanding\"]:\n",
    "    dataset = load_dataset(dataset_key)\n",
    "\n",
    "    # Get train completions\n",
    "    completion_data = load_completion_data(\"zs_cot\", dataset_key, \"text-davinci-002\")\n",
    "    \n",
    "    # Generate FT data\n",
    "    template = \"special\"\n",
    "    file_key = \"zs_cot_{}_{}_all\".format(template, dataset_key)\n",
    "    train, test = get_train_test_indices(dataset_key)\n",
    "    generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                           indices=train, include_incorrect=True)\n",
    "\n",
    "    # Create file on OpenAI\n",
    "    create_finetune_file(file_key)  # openai\n",
    "    \n",
    "    # Create finetunes\n",
    "    for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        model_key = \"{}_{}\".format(base_model, file_key)\n",
    "        create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bef10d91",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for dataset_key in [\"date_understanding\"]:\n",
    "        total += 1\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_all\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        if get_model_id(model_key, strict=False):\n",
    "            print(\"{:60s} exists\".format(model_key))\n",
    "            complete += 1\n",
    "            continue\n",
    "\n",
    "        finetune_id = get_finetune_id(model_key)\n",
    "        if finetune_id is None:\n",
    "            print(\"{:60s} no finetune found\".format(model_key))\n",
    "            continue\n",
    "\n",
    "        response = openai.FineTune.retrieve(finetune_id)\n",
    "        model_id = response[\"fine_tuned_model\"]\n",
    "        if model_id is not None:\n",
    "            set_model_id(model_key, model_id)\n",
    "            print(\"{:60s} fetched\".format(model_key))\n",
    "            print(\"    {:40s}\".format(model_id))\n",
    "            complete += 1\n",
    "        else:\n",
    "            print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "\n",
    "\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0759cd84",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot_long\"\n",
    "for dataset_key in [\"date_understanding\"]:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        print(\"#\" * 80)\n",
    "        print(dataset_key, base_model_key)\n",
    "        print(\"#\" * 80)\n",
    "        file_key = \"zs_cot_{}_{}_all\".format(template, dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        train, test = get_train_test_indices(dataset_key)\n",
    "        _ = infer_cot_completions(completion_key, dataset_key, model_key, indices=test,\n",
    "                                  template=template, max_tokens=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b31a7dfa",
   "metadata": {},
   "source": [
    "# Fine-tune-CoT (All Samples w/ Diverse Reasoning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5490e0fc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "datasets = [\"multiarith\", \"svamp\", \"date_understanding\", \"last_letter_concatenation\"]\n",
    "for augmentations in [1, 2, 4, 8]:\n",
    "    print(\" Generate {}aug Data \".format(augmentations).center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        print(dataset_key)\n",
    "        completion_key = \"zs_cot_aug\"\n",
    "        model_key = \"text-davinci-002\"\n",
    "        train_indices, _ = get_train_test_indices(dataset_key)\n",
    "        generate_cot_completions(completion_key, dataset_key, model_key, indices=train_indices,\n",
    "                                 augmentations=augmentations)\n",
    "\n",
    "    print(\" Generate File, Upload File, Run Fine-tune \".center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        completion_data = load_completion_data(\"zs_cot_aug\", dataset_key, \"text-davinci-002\")\n",
    "        train_indices, _ = get_train_test_indices(dataset_key)\n",
    "\n",
    "        # Generate FT data\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_all_{}aug\".format(template, dataset_key, augmentations)\n",
    "        _ = generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                                   indices=train_indices, augmentations=augmentations, include_incorrect=True)\n",
    "\n",
    "        # Create file on OpenAI\n",
    "        file_id = get_file_id(file_key)\n",
    "        if file_id:\n",
    "            print(\"{:60s}{}\".format(file_key, file_id))\n",
    "        else:\n",
    "            create_finetune_file(file_key)  # openai\n",
    "\n",
    "        # Create finetunes\n",
    "        for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            model_key = \"{}_{}\".format(base_model, file_key)\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id:\n",
    "                print(\"{:60s}{}\".format(model_key, finetune_id))\n",
    "            else:\n",
    "                create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45827f76",
   "metadata": {},
   "source": [
    "### Fetch Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db7a11a0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "\n",
    "# datasets = [\"multiarith\", \"svamp\"]\n",
    "datasets = [\"multiarith\", \"svamp\", \"date_understanding\", \"last_letter_concatenation\"]\n",
    "for augmentations in [1, 2, 4, 8]:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        for dataset_key in datasets:\n",
    "            total += 1\n",
    "            template = \"special\"\n",
    "            file_key = \"zs_cot_{}_{}_all_{}aug\".format(template, dataset_key, augmentations)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False):\n",
    "                print(\"{:60s} exists\".format(model_key))\n",
    "                complete += 1\n",
    "                continue\n",
    "\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id is None:\n",
    "                print(\"{:60s} no finetune found\".format(model_key))\n",
    "                continue\n",
    "\n",
    "            response = openai.FineTune.retrieve(finetune_id)\n",
    "            model_id = response[\"fine_tuned_model\"]\n",
    "            if model_id is not None:\n",
    "                set_model_id(model_key, model_id)\n",
    "                print(\"{:60s} fetched\".format(model_key))\n",
    "                print(\"    {:40s}\".format(model_id))\n",
    "                complete += 1\n",
    "            else:\n",
    "                print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bd459cc",
   "metadata": {},
   "source": [
    "### Long Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc706641",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# datasets = [\"multiarith\", \"svamp\"]\n",
    "datasets = [\"multiarith\", \"svamp\", \"date_understanding\", \"last_letter_concatenation\"]\n",
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot_long\"\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "#     for augmentations in [1, 2, 4, 8, 16, 32, 64]:\n",
    "    for augmentations in [1, 2, 4, 8]:\n",
    "        for dataset_key in datasets:\n",
    "            file_key = \"zs_cot_{}_{}_all_{}aug\".format(template, dataset_key, augmentations)\n",
    "            print(\" {} \".format(file_key).center(80, \"#\"))\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False) is not None:\n",
    "                _ = infer_cot_completions(completion_key, dataset_key, model_key,\n",
    "                                          template=template, max_tokens=1024)\n",
    "            else:\n",
    "                print(\"Model not found\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46ba419e",
   "metadata": {},
   "source": [
    "# Few-shot Fine-tune-CoT (Fine-tuning with Few-Shot Samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "082a145e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for shots in [8, 32, 128]:\n",
    "    print(\" Generate {}shot Data \".format(shots).center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        print(dataset_key)\n",
    "        completion_key = \"zs_cot\"\n",
    "        model_key = \"text-davinci-002\"\n",
    "        indices = get_few_shot_train_indices(dataset_key, shots=shots)\n",
    "        generate_cot_completions(completion_key, dataset_key, model_key, indices=indices)\n",
    "\n",
    "    print(\" Correct Samples \".center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        completion_data = load_completion_data(\"zs_cot\", dataset_key, \"text-davinci-002\")\n",
    "        indices = get_few_shot_train_indices(dataset_key, shots=shots)\n",
    "        evaluation = evaluate_completions(completion_data, dataset_key, template=None,\n",
    "                                          indices=indices)\n",
    "        correct = get_evaluation_metrics(evaluation)[\"accuracy\"] * shots\n",
    "        print(\"{:40s}: {:03d}\".format(dataset_key, round(correct)))\n",
    "\n",
    "\n",
    "    print(\" Generating FT Data and Registering File \".center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        completion_data = load_completion_data(\"zs_cot\", dataset_key, \"text-davinci-002\")\n",
    "        indices = get_few_shot_train_indices(dataset_key, shots=shots)\n",
    "\n",
    "        # Generate FT data\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_{}shot\".format(template, dataset_key, shots)\n",
    "        data = generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                                      indices=indices)\n",
    "\n",
    "        # Create file on OpenAI\n",
    "        file_id = get_file_id(file_key)\n",
    "        if file_id:\n",
    "            print(\"{:60s}{}\".format(file_key, file_id))\n",
    "        else:\n",
    "            create_finetune_file(file_key)  # openai\n",
    "\n",
    "        # Create finetunes\n",
    "        for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            model_key = \"{}_{}\".format(base_model, file_key)\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id:\n",
    "                print(\"{:60s}{}\".format(model_key, finetune_id))\n",
    "            else:\n",
    "                create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09fafcaa",
   "metadata": {},
   "source": [
    "### Fetch Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9513a37a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "\n",
    "for shots in [8, 32, 128]:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        for dataset_key in datasets:\n",
    "            total += 1\n",
    "            template = \"special\"\n",
    "            file_key = \"zs_cot_{}_{}_{}shot\".format(template, dataset_key, shots)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False):\n",
    "                print(\"{:60s} exists\".format(model_key))\n",
    "                complete += 1\n",
    "                continue\n",
    "\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id is None:\n",
    "                print(\"{:60s} no finetune found\".format(model_key))\n",
    "                continue\n",
    "\n",
    "            response = openai.FineTune.retrieve(finetune_id)\n",
    "            model_id = response[\"fine_tuned_model\"]\n",
    "            if model_id is not None:\n",
    "                set_model_id(model_key, model_id)\n",
    "                print(\"{:60s} fetched\".format(model_key))\n",
    "                print(\"    {:40s}\".format(model_id))\n",
    "                complete += 1\n",
    "            else:\n",
    "                print(\"{:60s} pending\".format(model_key))\n",
    "\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fea5178",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f589ffcc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for shots in [8, 32, 128]:\n",
    "    template = \"special\"\n",
    "    completion_key  = \"finetune_cot\"\n",
    "    for dataset_key in datasets:\n",
    "        for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            print(\" {}-{}-{}shot \".format(dataset_key, base_model_key, shots).center(80, \"-\"))\n",
    "            file_key = \"zs_cot_{}_{}_{}shot\".format(template, dataset_key, shots)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False) is not None:\n",
    "                _ = infer_cot_completions(completion_key, dataset_key, model_key, template=template)\n",
    "            else:\n",
    "                print(\"Model not found\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dd919ea",
   "metadata": {},
   "source": [
    "### Long Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7606a0ad",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for shots in [8, 32, 128]:\n",
    "    template = \"special\"\n",
    "    completion_key  = \"finetune_cot_long\"\n",
    "    for dataset_key in datasets:\n",
    "        for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            print(\" {}-{}-{}shot \".format(dataset_key, base_model_key, shots).center(80, \"-\"))\n",
    "            file_key = \"zs_cot_{}_{}_{}shot\".format(template, dataset_key, shots)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False) is not None:\n",
    "                _ = infer_cot_completions(completion_key, dataset_key, model_key,\n",
    "                                          template=template, max_tokens=1024)\n",
    "            else:\n",
    "                print(\"Model not found\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be96d690",
   "metadata": {},
   "source": [
    "# Vanilla FT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcb32033",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "datasets = ALL_DATASETS\n",
    "for dataset_key in datasets:\n",
    "    dataset = load_dataset(dataset_key)\n",
    "    train_indices, _ = get_train_test_indices(dataset_key)\n",
    "\n",
    "    # Generate FT data\n",
    "    file_key = \"{}_train\".format(dataset_key)\n",
    "    data = generate_finetune_data(dataset, dataset_key, template=None, file_key=file_key, indices=train_indices)\n",
    "\n",
    "    # Create file on OpenAI\n",
    "    file_id = get_file_id(file_key)\n",
    "    if file_id:\n",
    "        print(\"{:60s}{}\".format(file_key, file_id))\n",
    "    else:\n",
    "        create_finetune_file(file_key)  # openai\n",
    "\n",
    "    # Create finetunes\n",
    "    for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        model_key = \"{}_{}\".format(base_model, file_key)\n",
    "        finetune_id = get_finetune_id(model_key)\n",
    "        if finetune_id:\n",
    "            print(\"{:60s}{}\".format(model_key, finetune_id))\n",
    "        else:\n",
    "            create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e4d423f",
   "metadata": {},
   "source": [
    "### Fetch Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c71ad567",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "\n",
    "print(\"-\" * 80)\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for dataset_key in datasets:\n",
    "        total += 1\n",
    "        file_key = \"{}_train\".format(dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        if get_model_id(model_key, strict=False):\n",
    "            print(\"{:60s} exists\".format(model_key))\n",
    "            complete += 1\n",
    "            continue\n",
    "\n",
    "        finetune_id = get_finetune_id(model_key)\n",
    "        if finetune_id is None:\n",
    "            print(\"{:60s} no finetune found\".format(model_key))\n",
    "            continue\n",
    "\n",
    "        response = openai.FineTune.retrieve(finetune_id)\n",
    "        model_id = response[\"fine_tuned_model\"]\n",
    "        if model_id is not None:\n",
    "            set_model_id(model_key, model_id)\n",
    "            print(\"{:60s} fetched\".format(model_key))\n",
    "            print(\"    {:40s}\".format(model_id))\n",
    "            complete += 1\n",
    "        else:\n",
    "            print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "            \n",
    "print(\"-\" * 80)\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70c632cb",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "822c9a9a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "completion_key  = \"ft\"\n",
    "for dataset_key in datasets:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        print(\" {}-{} \".format(dataset_key, base_model_key).center(80, \"-\"))\n",
    "        file_key = \"{}_train\".format(dataset_key)\n",
    "        model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "        if get_model_id(model_key, strict=False) is not None:\n",
    "            completion_data = infer_cot_completions(completion_key, dataset_key, model_key, template=None)\n",
    "        else:\n",
    "            print(\"Model not found\")\n",
    "        print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bd6c1f8",
   "metadata": {},
   "source": [
    "# Few-Shot Vanilla FT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7364e3d0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "datasets = ALL_DATASETS\n",
    "for shots in [8, 32, 128]:\n",
    "    print(\" Generate {}shot Data \".format(shots).center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        dataset = load_dataset(dataset_key)\n",
    "        indices = get_few_shot_train_indices(dataset_key, shots=shots)\n",
    "\n",
    "        # Generate FT data\n",
    "        file_key = \"{}_{}shot\".format(dataset_key, shots)\n",
    "        data = generate_finetune_data(dataset, dataset_key, template=None, file_key=file_key, indices=indices)\n",
    "\n",
    "        # Create file on OpenAI\n",
    "        file_id = get_file_id(file_key)\n",
    "        if file_id:\n",
    "            print(\"{:60s}{}\".format(file_key, file_id))\n",
    "        else:\n",
    "            create_finetune_file(file_key)  # openai\n",
    "\n",
    "        # Create finetunes\n",
    "        for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            model_key = \"{}_{}\".format(base_model, file_key)\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id:\n",
    "                print(\"{:60s}{}\".format(model_key, finetune_id))\n",
    "            else:\n",
    "                create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b85144f",
   "metadata": {},
   "source": [
    "### Fetch Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e39627d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "\n",
    "print(\"-\" * 80)\n",
    "for shots in [8, 32, 128]:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        for dataset_key in datasets:\n",
    "            total += 1\n",
    "            file_key = \"{}_{}shot\".format(dataset_key, shots)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False):\n",
    "                print(\"{:60s} exists\".format(model_key))\n",
    "                complete += 1\n",
    "                continue\n",
    "\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id is None:\n",
    "                print(\"{:60s} no finetune found\".format(model_key))\n",
    "                continue\n",
    "\n",
    "            response = openai.FineTune.retrieve(finetune_id)\n",
    "            model_id = response[\"fine_tuned_model\"]\n",
    "            if model_id is not None:\n",
    "                set_model_id(model_key, model_id)\n",
    "                print(\"{:60s} fetched\".format(model_key))\n",
    "                print(\"    {:40s}\".format(model_id))\n",
    "                complete += 1\n",
    "            else:\n",
    "                print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "            \n",
    "print(\"-\" * 80)\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "922faf9d",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a2968b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "completion_key  = \"ft\"\n",
    "for shots in [8, 32, 128]:\n",
    "    for dataset_key in datasets:\n",
    "        for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            print(\" {}-{} \".format(dataset_key, base_model_key).center(80, \"-\"))\n",
    "            file_key = \"{}_{}shot\".format(dataset_key, shots)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False) is not None:\n",
    "                completion_data = infer_cot_completions(completion_key, dataset_key, model_key, template=None)\n",
    "            else:\n",
    "                print(\"Model not found\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44101773",
   "metadata": {},
   "source": [
    "## \\* Diverse Reasoning for Confidence Sampling (Date)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d8831eb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for augmentations in [1, 2, 4, 8, 16]:\n",
    "    print(\" Generate {}aug Data \".format(augmentations).center(80, \"#\"))\n",
    "    dataset_key = \"date_understanding\"\n",
    "    completion_key = \"zs_cot_aug\"\n",
    "    model_key = \"text-davinci-002\"\n",
    "    train_indices, _ = get_train_test_indices(dataset_key)\n",
    "    generate_cot_completions(completion_key, dataset_key, model_key, indices=train_indices,\n",
    "                             augmentations=augmentations, reasoning_temperature=0.7)\n",
    "\n",
    "    print(\" Augmented Accuracy (Train) \".center(80, \"#\"))\n",
    "    completion_data = load_completion_data(\"zs_cot_aug\", dataset_key, \"text-davinci-002\")\n",
    "    train_indices, _ = get_train_test_indices(dataset_key)\n",
    "    evaluation = evaluate_completions(completion_data, dataset_key, template=None,\n",
    "                                      indices=train_indices, augmentations=augmentations)\n",
    "    accuracy = get_evaluation_metrics(evaluation)[\"accuracy\"]\n",
    "    print(\"{:60s}: {:f}\".format(dataset_key, accuracy))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e31fa951",
   "metadata": {},
   "source": [
    "# Fine-tune-CoT (Diverse Reasoning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "176a2de1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "datasets = [\"multiarith\", \"svamp\"]\n",
    "datasets = [\"date_understanding\", \"last_letter_concatenation\"]\n",
    "datasets = [\"single_eq\", \"addsub\", \"tracking_shuffled_objects\", \"coin_flip\", \"strategy_qa\"]\n",
    "for augmentations in [2, 1, 4, 8]:\n",
    "    print(\" Generate {}aug Data \".format(augmentations).center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        print(dataset_key)\n",
    "        completion_key = \"zs_cot_aug\"\n",
    "        model_key = \"text-davinci-002\"\n",
    "        train_indices, _ = get_train_test_indices(dataset_key)\n",
    "        generate_cot_completions(completion_key, dataset_key, model_key, indices=train_indices,\n",
    "                                 augmentations=augmentations)\n",
    "\n",
    "    print(\" Augmented Accuracy (Train) \".center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        completion_data = load_completion_data(\"zs_cot_aug\", dataset_key, \"text-davinci-002\")\n",
    "        train_indices, _ = get_train_test_indices(dataset_key)\n",
    "        evaluation = evaluate_completions(completion_data, dataset_key, template=None,\n",
    "                                          indices=train_indices, augmentations=augmentations)\n",
    "        accuracy = get_evaluation_metrics(evaluation)[\"accuracy\"]\n",
    "        print(\"{:60s}: {:f}\".format(dataset_key, accuracy))\n",
    "\n",
    "\n",
    "    print(\" Generate File, Upload File, Run Fine-tune \".center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        completion_data = load_completion_data(\"zs_cot_aug\", dataset_key, \"text-davinci-002\")\n",
    "        train_indices, _ = get_train_test_indices(dataset_key)\n",
    "\n",
    "        # Generate FT data\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_{}aug\".format(template, dataset_key, augmentations)\n",
    "        _ = generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                                   indices=train_indices, augmentations=augmentations)\n",
    "\n",
    "        # Create file on OpenAI\n",
    "        file_id = get_file_id(file_key)\n",
    "        if file_id:\n",
    "            print(\"{:60s}{}\".format(file_key, file_id))\n",
    "        else:\n",
    "            create_finetune_file(file_key)  # openai\n",
    "\n",
    "        # Create finetunes\n",
    "        for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            model_key = \"{}_{}\".format(base_model, file_key)\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id:\n",
    "                print(\"{:60s}{}\".format(model_key, finetune_id))\n",
    "            else:\n",
    "                create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "693068e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = [\"svamp\"]\n",
    "for augmentations in [64]:\n",
    "    print(\" Generate {}aug Data \".format(augmentations).center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        print(dataset_key)\n",
    "        completion_key = \"zs_cot_aug\"\n",
    "        model_key = \"text-davinci-002\"\n",
    "        train_indices, _ = get_train_test_indices(dataset_key)\n",
    "        generate_cot_completions(completion_key, dataset_key, model_key, indices=train_indices,\n",
    "                                 augmentations=augmentations)\n",
    "\n",
    "    print(\" Augmented Accuracy (Train) \".center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        completion_data = load_completion_data(\"zs_cot_aug\", dataset_key, \"text-davinci-002\")\n",
    "        train_indices, _ = get_train_test_indices(dataset_key)\n",
    "        evaluation = evaluate_completions(completion_data, dataset_key, template=None,\n",
    "                                          indices=train_indices, augmentations=augmentations)\n",
    "        accuracy = get_evaluation_metrics(evaluation)[\"accuracy\"]\n",
    "        print(\"{:60s}: {:f}\".format(dataset_key, accuracy))\n",
    "\n",
    "\n",
    "    print(\" Generate File, Upload File, Run Fine-tune \".center(80, \"#\"))\n",
    "    for dataset_key in datasets:\n",
    "        completion_data = load_completion_data(\"zs_cot_aug\", dataset_key, \"text-davinci-002\")\n",
    "        train_indices, _ = get_train_test_indices(dataset_key)\n",
    "\n",
    "        # Generate FT data\n",
    "        template = \"special\"\n",
    "        file_key = \"zs_cot_{}_{}_{}aug\".format(template, dataset_key, augmentations)\n",
    "        _ = generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                                   indices=train_indices, augmentations=augmentations)\n",
    "\n",
    "        # Create file on OpenAI\n",
    "        file_id = get_file_id(file_key)\n",
    "        if file_id:\n",
    "            print(\"{:60s}{}\".format(file_key, file_id))\n",
    "        else:\n",
    "            create_finetune_file(file_key)  # openai\n",
    "\n",
    "        # Create finetunes\n",
    "        for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            model_key = \"{}_{}\".format(base_model, file_key)\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id:\n",
    "                print(\"{:60s}{}\".format(model_key, finetune_id))\n",
    "            else:\n",
    "                create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e26dbc0d",
   "metadata": {},
   "source": [
    "### Fetch Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df1fda6c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "\n",
    "# datasets = [\"multiarith\", \"svamp\"]\n",
    "# datasets = [\"date_understanding\", \"last_letter_concatenation\"]\n",
    "datasets = [\"single_eq\", \"addsub\", \"tracking_shuffled_objects\", \"coin_flip\", \"strategy_qa\"]\n",
    "for augmentations in [1, 2, 4, 8]:\n",
    "    for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        for dataset_key in datasets:\n",
    "            total += 1\n",
    "            template = \"special\"\n",
    "            file_key = \"zs_cot_{}_{}_{}aug\".format(template, dataset_key, augmentations)\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False):\n",
    "                print(\"{:60s} exists\".format(model_key))\n",
    "                complete += 1\n",
    "                continue\n",
    "\n",
    "            finetune_id = get_finetune_id(model_key)\n",
    "            if finetune_id is None:\n",
    "                print(\"{:60s} no finetune found\".format(model_key))\n",
    "                continue\n",
    "\n",
    "            response = openai.FineTune.retrieve(finetune_id)\n",
    "            model_id = response[\"fine_tuned_model\"]\n",
    "            if model_id is not None:\n",
    "                set_model_id(model_key, model_id)\n",
    "                print(\"{:60s} fetched\".format(model_key))\n",
    "                print(\"    {:40s}\".format(model_id))\n",
    "                complete += 1\n",
    "            else:\n",
    "                print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28685574",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46601abe",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# datasets = [\"multiarith\", \"svamp\"]\n",
    "datasets = [\"date_understanding\", \"last_letter_concatenation\"]\n",
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot\"\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "#     for augmentations in [1, 2, 4, 8, 16, 32, 64]:\n",
    "    for augmentations in [1, 2, 4, 8]:\n",
    "        for dataset_key in datasets:\n",
    "            file_key = \"zs_cot_{}_{}_{}aug\".format(template, dataset_key, augmentations)\n",
    "            print(\" {} \".format(file_key).center(80, \"#\"))\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False) is not None:\n",
    "                _ = infer_cot_completions(completion_key, dataset_key, model_key, template=template)\n",
    "            else:\n",
    "                print(\"Model not found\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fc914b2",
   "metadata": {},
   "source": [
    "### Long Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42b4213a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# datasets = [\"multiarith\", \"svamp\"]\n",
    "datasets = [\"date_understanding\", \"last_letter_concatenation\"]\n",
    "datasets = [\"single_eq\", \"addsub\", \"tracking_shuffled_objects\", \"coin_flip\", \"strategy_qa\"]\n",
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot_long\"\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "#     for augmentations in [1, 2, 4, 8, 16, 32, 64]:\n",
    "    for augmentations in [8]:\n",
    "        for dataset_key in datasets:\n",
    "            file_key = \"zs_cot_{}_{}_{}aug\".format(template, dataset_key, augmentations)\n",
    "            print(\" {} \".format(file_key).center(80, \"#\"))\n",
    "            model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "            if get_model_id(model_key, strict=False) is not None:\n",
    "                _ = infer_cot_completions(completion_key, dataset_key, model_key,\n",
    "                                          template=template, max_tokens=1024)\n",
    "            else:\n",
    "                print(\"Model not found\")\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b95067b",
   "metadata": {},
   "source": [
    "# Few-shot Fine-tune-CoT (Diverse Reasoning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16b52701",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "datasets = [\"multiarith\", \"svamp\"]\n",
    "for shots in [8, 32, 128]:\n",
    "    for augmentations in [1, 2, 4, 8, 16, 32, 64]:\n",
    "        print(\" Generate {}shot {}aug Data \".format(shots, augmentations).center(80, \"#\"))\n",
    "        for dataset_key in datasets:\n",
    "            print(dataset_key)\n",
    "            completion_key = \"zs_cot_aug\"\n",
    "            model_key = \"text-davinci-002\"\n",
    "            indices = get_few_shot_train_indices(dataset_key, shots=shots)\n",
    "            generate_cot_completions(completion_key, dataset_key, model_key, indices=indices,\n",
    "                                     augmentations=augmentations)\n",
    "\n",
    "        print(\" Augmented Accuracy (Train) \".center(80, \"#\"))\n",
    "        for dataset_key in datasets:\n",
    "            completion_data = load_completion_data(\"zs_cot_aug\", dataset_key, \"text-davinci-002\")\n",
    "            indices = get_few_shot_train_indices(dataset_key, shots=shots)\n",
    "            evaluation = evaluate_completions(completion_data, dataset_key, template=None,\n",
    "                                              indices=indices, augmentations=augmentations)\n",
    "            accuracy = get_evaluation_metrics(evaluation)[\"accuracy\"]\n",
    "            print(\"{:60s}: {:f}\".format(dataset_key, accuracy))\n",
    "\n",
    "        print(\" Generate File, Upload File, Run Fine-tune \".center(80, \"#\"))\n",
    "        for dataset_key in datasets:\n",
    "            completion_data = load_completion_data(\"zs_cot_aug\", dataset_key, \"text-davinci-002\")\n",
    "            indices = get_few_shot_train_indices(dataset_key, shots=shots)\n",
    "\n",
    "            # Generate FT data\n",
    "            template = \"special\"\n",
    "            file_key = \"zs_cot_{}_{}_{}shot_{}aug\".format(template, dataset_key, shots, augmentations)\n",
    "            _ = generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                                       indices=indices, augmentations=augmentations)\n",
    "\n",
    "            # Create file on OpenAI\n",
    "            file_id = get_file_id(file_key)\n",
    "            if file_id:\n",
    "                print(\"{:60s}{}\".format(file_key, file_id))\n",
    "            else:\n",
    "                create_finetune_file(file_key)  # openai\n",
    "\n",
    "            # Create finetunes\n",
    "            for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "                model_key = \"{}_{}\".format(base_model, file_key)\n",
    "                finetune_id = get_finetune_id(model_key)\n",
    "                if finetune_id:\n",
    "                    print(\"{:60s}{}\".format(model_key, finetune_id))\n",
    "                else:\n",
    "                    create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3da8a0c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "total = 0\n",
    "complete = 0\n",
    "\n",
    "datasets = [\"multiarith\", \"svamp\"]\n",
    "for shots in [8, 32, 128]:\n",
    "    for augmentations in [1, 2, 4, 8, 16, 32, 64]:\n",
    "        for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "            for dataset_key in datasets:\n",
    "                total += 1\n",
    "                template = \"special\"\n",
    "                file_key = \"zs_cot_{}_{}_{}shot_{}aug\".format(template, dataset_key, shots, augmentations)\n",
    "                model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "                if get_model_id(model_key, strict=False):\n",
    "                    print(\"{:60s} exists\".format(model_key))\n",
    "                    complete += 1\n",
    "                    continue\n",
    "\n",
    "                finetune_id = get_finetune_id(model_key)\n",
    "                if finetune_id is None:\n",
    "                    print(\"{:60s} no finetune found\".format(model_key))\n",
    "                    continue\n",
    "\n",
    "                response = openai.FineTune.retrieve(finetune_id)\n",
    "                model_id = response[\"fine_tuned_model\"]\n",
    "                if model_id is not None:\n",
    "                    set_model_id(model_key, model_id)\n",
    "                    print(\"{:60s} fetched\".format(model_key))\n",
    "                    print(\"    {:40s}\".format(model_id))\n",
    "                    complete += 1\n",
    "                else:\n",
    "                    print(\"{:60s} {}\".format(model_key, response[\"status\"]))\n",
    "\n",
    "print(\"{} of {} models are ready\".format(complete, total))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63bf30d4",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "870a8201",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "summary = []\n",
    "datasets = [\"multiarith\", \"svamp\"]\n",
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot\"\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for shots in [8, 32, 128]:\n",
    "        for augmentations in [1, 2, 4, 8, 16, 32, 64]:\n",
    "            for dataset_key in datasets:\n",
    "                file_key = \"zs_cot_{}_{}_{}shot_{}aug\".format(template, dataset_key, shots, augmentations)\n",
    "                print(\" {} \".format(file_key).center(80, \"#\"))\n",
    "                model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "                if get_model_id(model_key, strict=False) is not None:\n",
    "                    _ = infer_cot_completions(completion_key, dataset_key, model_key, template=template)\n",
    "                else:\n",
    "                    print(\"Model not found\")\n",
    "                print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ed9252f",
   "metadata": {},
   "source": [
    "### Long Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cc6af8f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "summary = []\n",
    "datasets = [\"multiarith\", \"svamp\"]\n",
    "template = \"special\"\n",
    "completion_key  = \"finetune_cot_long\"\n",
    "for base_model_key in [\"ada\", \"babbage\", \"curie\"]:\n",
    "    for shots in [8, 32, 128]:\n",
    "        for augmentations in [1, 2, 4, 8, 16, 32, 64]:\n",
    "            for dataset_key in datasets:\n",
    "                file_key = \"zs_cot_{}_{}_{}shot_{}aug\".format(template, dataset_key, shots, augmentations)\n",
    "                print(\" {} \".format(file_key).center(80, \"#\"))\n",
    "                model_key = \"{}_{}\".format(base_model_key, file_key)\n",
    "                if get_model_id(model_key, strict=False) is not None:\n",
    "                    _ = infer_cot_completions(completion_key, dataset_key, model_key,\n",
    "                                              template=template, max_tokens=1024)\n",
    "                else:\n",
    "                    print(\"Model not found\")\n",
    "                print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a34b54d",
   "metadata": {},
   "source": [
    "# Reasoning Length Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23f27918",
   "metadata": {},
   "outputs": [],
   "source": [
    "long_datasets = \"aqua,commonsense_qa,strategy_qa,gsm8k,svamp\".split(\",\")\n",
    "long_datasets_go = \"strategy_qa,gsm8k,svamp\".split(\",\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "350230c3",
   "metadata": {},
   "source": [
    "## Test Generation\n",
    "\n",
    "Generate test answers w/ full reason (2048) for analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876eed38",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "completion_key = \"zs_cot_long\"\n",
    "model_key = \"text-davinci-002\"\n",
    "for dataset_key in datasets:\n",
    "    train, test = get_train_test_indices(dataset_key)\n",
    "    generate_cot_completions(completion_key, dataset_key, model_key, indices=test,\n",
    "                             max_tokens=2048, request_batch_size=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c4020c8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_indices, test_indices = get_train_test_indices(\"gsm8k\")\n",
    "print(len(train_indices))\n",
    "train_indices, test_indices = get_train_test_indices(\"multiarith\")\n",
    "print(len(train_indices))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57530c65",
   "metadata": {},
   "source": [
    "### Test Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "442d7b23",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "metrics_by_dataset = dict()\n",
    "evaluation_by_dataset = dict()\n",
    "for dataset_key in long_datasets:\n",
    "    completion_data = load_completion_data(\"zs_cot_long\", dataset_key, \"text-davinci-002\")\n",
    "    train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "    evaluation = evaluate_completions(completion_data, dataset_key, template=None,\n",
    "                                      indices=test_indices)\n",
    "    evaluation_by_dataset[dataset_key] = evaluation\n",
    "    metrics_by_dataset[dataset_key] = get_evaluation_metrics(evaluation)\n",
    "with pd.option_context('display.float_format', '{:,.3f}'.format):\n",
    "    display(pd.DataFrame(metrics_by_dataset).T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23d99b04",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from data.tokens import get_token_count\n",
    "\n",
    "lengths_by_dataset = defaultdict(dict)\n",
    "for dataset_key in long_datasets:\n",
    "    average = 0\n",
    "    print(\" {} \".format(dataset_key).center(80, \"#\"))\n",
    "    train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "    samples = load_completion_data(completion_key, dataset_key, model_key)\n",
    "    for i in test_indices:\n",
    "        s = samples[i][0]\n",
    "        lengths_by_dataset[dataset_key][i] = get_token_count(s[\"reasoning_completion\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5a26634",
   "metadata": {},
   "outputs": [],
   "source": [
    "thresholds = [0, 64, 128, 256, 384, 512, 768, 1024, 2047, 2048]\n",
    "accuracy_by_threshold_by_dataset = dict()\n",
    "for dataset in long_datasets:\n",
    "    lengths = lengths_by_dataset[dataset]\n",
    "    e = evaluation_by_dataset[dataset].set_index(\"sample_index\")\n",
    "    l = pd.Series(lengths)\n",
    "    l.name = \"length\"\n",
    "    df = pd.concat([e, l], axis=1)\n",
    "    df = df.loc[:, [\"correct\", \"length\"]]\n",
    "    accuracy_by_threshold = defaultdict(list)\n",
    "    for t1, t2 in zip(thresholds[:-1], thresholds[1:]):\n",
    "        samples = df[(df.length > t1) & (df.length <= t2)]\n",
    "        accuracy = samples.correct.sum() / samples.correct.count()\n",
    "        accuracy_by_threshold[\"t1\"].append(t1)\n",
    "        accuracy_by_threshold[\"t2\"].append(t2)\n",
    "        accuracy_by_threshold[\"count\"].append(samples.correct.count())\n",
    "        accuracy_by_threshold[\"correct\"].append(samples.correct.sum())\n",
    "        accuracy_by_threshold[\"accurcy\"].append(accuracy)\n",
    "    accuracy_by_threshold = pd.DataFrame(accuracy_by_threshold)\n",
    "    accuracy_by_threshold_by_dataset[dataset] = accuracy_by_threshold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16b92fcc",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "accuracy_by_threshold_by_dataset[\"strategy_qa\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba5dadb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "accuracy_by_threshold_by_dataset[\"commonsense_qa\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b12ea6b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "accuracy_by_threshold_by_dataset[\"gsm8k\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49155b42",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "accuracy_by_threshold_by_dataset[\"svamp\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "662ef524",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "accuracy_by_threshold_by_dataset[\"aqua\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe3e1afe",
   "metadata": {},
   "outputs": [],
   "source": [
    "long_reason_indices = defaultdict(list)\n",
    "inf_reason_indices = defaultdict(list)\n",
    "\n",
    "for dataset in long_datasets_go:\n",
    "    lengths = lengths_by_dataset[dataset]\n",
    "    for i, length in lengths.items():\n",
    "        if length == 2048:\n",
    "            inf_reason_indices[dataset].append(i)\n",
    "        elif length > 512:\n",
    "            long_reason_indices[dataset].append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35547380",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for dataset in long_datasets_go:\n",
    "    print(dataset.upper().center(80, \"#\"))\n",
    "    completion_key = \"zs_cot_long\"\n",
    "    model_key = \"text-davinci-002\"\n",
    "    samples = load_completion_data(completion_key, dataset, model_key)\n",
    "    for i in inf_reason_indices[dataset]:\n",
    "        s = samples[i]\n",
    "        print()\n",
    "        print(\"#\" * 80)\n",
    "        print(\"SAMPLE #{}\".format(s[0][\"sample_index\"]))\n",
    "        print(s[0][\"question\"])\n",
    "        print(\"-\" * 80)\n",
    "        print(s[0][\"reasoning_completion\"])\n",
    "        print(\"-\" * 80)\n",
    "        print(\"Answer\", s[0][\"answer\"])\n",
    "        print(\"-\" * 80)\n",
    "        print(\"Tokens\", get_token_count(s[0][\"reasoning_completion\"]))\n",
    "        print(\"-\" * 80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "863afb1f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for dataset in long_datasets_go:\n",
    "    print(dataset.upper().center(80, \"#\"))\n",
    "    completion_key = \"zs_cot_long\"\n",
    "    model_key = \"text-davinci-002\"\n",
    "    samples = load_completion_data(completion_key, dataset, model_key)\n",
    "    for i in long_reason_indices[dataset]:\n",
    "        s = samples[i]\n",
    "        print()\n",
    "        print(\"#\" * 80)\n",
    "        print(\"SAMPLE #{}\".format(s[0][\"sample_index\"]))\n",
    "        print(s[0][\"question\"])\n",
    "        print(\"-\" * 80)\n",
    "        print(s[0][\"reasoning_completion\"])\n",
    "        print(\"-\" * 80)\n",
    "        print(\"Answer\", s[0][\"answer\"])\n",
    "        print(\"-\" * 80)\n",
    "        print(\"Tokens\", get_token_count(s[0][\"reasoning_completion\"]))\n",
    "        print(\"-\" * 80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edcd7e5d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from data.tokens import get_token_count\n",
    "\n",
    "for dataset_key in long_datasets:\n",
    "    print(\"#\" * 80)\n",
    "    print(dataset_key)\n",
    "    print(\"#\" * 80)\n",
    "    samples = load_completion_data(completion_key, dataset_key, model_key)\n",
    "    for s in samples:\n",
    "        if get_token_count(s[0][\"reasoning_completion\"]) > 512:\n",
    "            print(s[0][\"question\"])\n",
    "            print(\"-\" * 80)\n",
    "            print(s[0][\"reasoning_completion\"])\n",
    "            print(\"-\" * 80)\n",
    "            print(\"Answer\", s[0][\"answer\"])\n",
    "            print(\"-\" * 80)\n",
    "            print(\"Tokens\", get_token_count(s[0][\"reasoning_completion\"]))\n",
    "            print(\"-\" * 80)\n",
    "            break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ebc46af",
   "metadata": {},
   "source": [
    "## Train Generation\n",
    "\n",
    "Note, train answers are generated w/ truncated reason (512) for fine-tuning, test answers have been generated w/ full reason for analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db9b4481",
   "metadata": {},
   "outputs": [],
   "source": [
    "long_datasets = \"aqua,commonsense_qa,strategy_qa,gsm8k,svamp\".split(\",\")\n",
    "# long_datasets = \"strategy_qa,gsm8k,svamp\".split(\",\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c6187da",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "completion_key = \"zs_cot_long\"\n",
    "model_key = \"text-davinci-002\"\n",
    "for dataset_key in long_datasets:\n",
    "    print(dataset_key.upper().center(80, \"=\"))\n",
    "    train, test = get_train_test_indices(dataset_key)\n",
    "    generate_cot_completions(completion_key, dataset_key, model_key, indices=train,\n",
    "                             max_tokens=2048, max_tokens_used=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff6a2a4b",
   "metadata": {},
   "source": [
    "# Curate and Fine-Tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d31e00e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for dataset_key in long_datasets:\n",
    "    dataset = load_dataset(dataset_key)\n",
    "    train_indices, test_indices = get_train_test_indices(dataset_key)\n",
    "\n",
    "    # Get train completions\n",
    "    completion_data = load_completion_data(\"zs_cot_long\", dataset_key, \"text-davinci-002\")\n",
    "    train_completions = {i: completion_data[i] for i in train_indices}\n",
    "    test_completions = {i: completion_data[i] for i in test_indices}\n",
    "\n",
    "    # Generate FT data\n",
    "    template = \"special\"\n",
    "    file_key = \"zs_cot_long_{}_{}_train\".format(template, dataset_key)\n",
    "    generate_finetune_data(completion_data, dataset_key, template=template, file_key=file_key,\n",
    "                          indices=train_indices)\n",
    "\n",
    "    # Create file on OpenAI\n",
    "    create_finetune_file(file_key)  # openai\n",
    "    \n",
    "    # Create finetunes\n",
    "    for base_model in [\"ada\", \"babbage\", \"curie\"]:\n",
    "        model_key = \"{}_{}\".format(base_model, file_key)\n",
    "        create_finetune(file_key, model_key=model_key, model=base_model, suffix=file_key[:40])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
